import asyncio
import cv2
import numpy as np
import av
import uvicorn
from fastapi import FastAPI, HTTPException, status
from fastapi.responses import HTMLResponse, StreamingResponse

from logger import log
from settings import settings

# 在这里配置你的所有RTSP视频流
# STREAMS_CONFIG = [
#     {
#         "id": 1,
#         "name": "测试视频流1 (文件)",
#         # 注意：请将这里的URL替换成你自己的有效RTSP地址
#         "rtsp": "rtsp://admin:admin@192.168.110.100:8557/h264"
#     },
#     {
#         "id": 2,
#         "name": "测试视频流2 (另一个文件)",
#         "rtsp": "rtsp://wowzaec2demo.streamlock.net/vod/mp4:BigBuckBunny_115k.mov"
#     },
#     # 你可以继续添加更多视频流
#     # {
#     #     "id": 3,
#     #     "name": "摄像头3",
#     #     "rtsp": "rtsp://user:password@your_camera_ip:554/stream1"
#     # },
# ]

# --- 应用核心逻辑 ---

app = FastAPI()

# 全局变量，用于存储所有视频流的状态
# 结构:
# {
#     stream_id: {
#         "name": "Stream Name",
#         "rtsp": "rtsp://...",
#         "frame": np.ndarray | None,  # 最新一帧图像
#         "count": 0,                   # 当前观看者数量
#         "lock": asyncio.Lock(),      # 访问此字典的锁
#         "task": asyncio.Task | None  # 拉流的后台任务
#     }
# }
VIDEO_STREAMS = {}


async def read_rtsp_stream(stream_id: int):
    """
    一个在后台运行的任务，负责从RTSP源读取帧并更新到全局VIDEO_STREAMS中。
    """
    stream_info = VIDEO_STREAMS[stream_id]
    rtsp = stream_info["rtsp"]
    log.info(f"开始拉取视频流: ID={stream_id}, URL={rtsp}")

    try:
        # 使用pyav打开RTSP流，设置超时和TCP传输
        container = av.open(rtsp,
                            format='rtsp',
                            options={'rtsp_transport': 'tcp',
                                     'framerate': '20',
                                     'video_size': '640x360'},
                            timeout=10)

        stream = container.streams.video[0]

        for frame in container.decode(stream):
            # 将帧转换为numpy数组 (BGR格式)
            img = frame.to_ndarray(format='bgr24')

            # 异步安全地更新帧
            async with stream_info["lock"]:
                stream_info["frame"] = img

            # 稍微等待，让出CPU给其他协程，并控制帧率
            await asyncio.sleep(0.02)  # 大约 50fps，可以根据需要调整

    except av.error.ExitError:
        log.warning(f"RTSP流读取正常结束: ID={stream_id}")
    except asyncio.CancelledError:
        log.info(f"拉流任务被取消: ID={stream_id}")
    except Exception as e:
        log.error(f"打开或读取RTSP流时发生错误: ID={stream_id}, Error: {e}")
        # 创建一个错误提示图片作为最后一帧
        error_frame = np.zeros((360, 640, 3), dtype=np.uint8)
        cv2.putText(error_frame, 'Stream Error', (180, 180),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        async with stream_info["lock"]:
            stream_info["frame"] = error_frame
    finally:
        log.info(f"停止拉取视频流: ID={stream_id}")
        if 'container' in locals():
            container.close()


async def frame_generator(stream_id: int):
    """
    一个生成器，持续地从VIDEO_STREAMS获取最新帧并编码为JPEG格式，
    用于HTTP流式响应。
    """
    while True:
        async with VIDEO_STREAMS[stream_id]["lock"]:
            frame = VIDEO_STREAMS[stream_id]["frame"]

        if frame is not None:
            # 调整图像大小，编码为JPEG
            resized_frame = cv2.resize(frame, (640, 360))
            ret, buffer = cv2.imencode('.jpg', resized_frame)
            if not ret:
                continue
            frame_bytes = buffer.tobytes()
            # 产生给HTTP响应的帧数据
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')

        # 等待一小段时间，控制发送给浏览器的帧率
        await asyncio.sleep(0.2)  # 大约 5fps


@app.on_event("startup")
async def startup_event():
    """
    应用启动时，根据配置初始化VIDEO_STREAMS字典。
    """
    log.info("应用启动，正在初始化视频流配置...")
    for config in settings.streams:
        stream_id = config["id"]
        VIDEO_STREAMS[stream_id] = {
            "name": config["name"],
            "rtsp": config["rtsp"],
            "frame": None,
            "count": 0,
            "lock": asyncio.Lock(),
            "task": None,
        }
    log.info(f"初始化完成，共配置了 {len(VIDEO_STREAMS)} 个视频流。")


@app.get("/", response_class=HTMLResponse)
async def index():
    """
    主页，显示所有可用的视频流链接。
    """
    links = []
    for stream_id, stream_info in VIDEO_STREAMS.items():
        links.append(f'<li><a href="/video/stream/{stream_id}">{stream_info["name"]} (ID: {stream_id}:{stream_info["count"]})</a></li>')
    links_html = '<ul>' + ''.join(links) + '</ul>'
    return f'<html><head><title>视频流服务</title></head><body><h1>可用视频流</h1>{links_html}</body></html>'


@app.get("/{edge_code}/{stream_id}")
async def video_feed(edge_code: str, stream_id: int):
    """
    处理单个视频流的Web请求。
    """
    if stream_id not in VIDEO_STREAMS:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="视频流ID未找到")

    stream_info = VIDEO_STREAMS[stream_id]

    async def stream_lifecycle_manager():
        """
        管理拉流任务的生命周期，在生成器退出时执行清理。
        """
        try:
            # 这是一个异步生成器，会持续产生帧数据
            async for frame in frame_generator(stream_id):
                yield frame
        except asyncio.CancelledError:
            log.info(f"Web客户端断开连接: ID={stream_id}")
        finally:
            # 当客户端断开连接后，执行这里的清理逻辑
            async with stream_info["lock"]:
                stream_info["count"] -= 1
                log.info(f"一个观看者离开，流 {stream_id} 剩余观看者: {stream_info['count']}")
                # 如果是最后一个观看者，则取消后台拉流任务
                if stream_info["count"] == 0 and stream_info["task"] is not None:
                    log.info(f"没有观看者了，正在停止流 {stream_id} 的后台拉取任务...")
                    stream_info["task"].cancel()
                    stream_info["task"] = None

    # 有新的客户端连接
    async with stream_info["lock"]:
        stream_info["count"] += 1
        log.info(f"新的观看者连接，流 {stream_id} 当前观看者: {stream_info['count']}")
        # 如果这是第一个观看者，并且后台任务没有运行，则启动它
        if stream_info["count"] == 1 and (stream_info["task"] is None or stream_info["task"].done()):
            stream_info["task"] = asyncio.create_task(read_rtsp_stream(stream_id))

    return StreamingResponse(stream_lifecycle_manager(), media_type='multipart/x-mixed-replace; boundary=frame')

def start():
    # 使用 uvicorn 运行 FastAPI 应用
    # host="0.0.0.0" 表示监听所有网络接口，允许局域网内其他设备访问
    uvicorn.run(app, host="0.0.0.0", port=8080)


if __name__ == "__main__":
    start()
